from zhipuai import ZhipuAI
#from openai import OpenAI
from dotenv import load_dotenv, find_dotenv

class Task:
    def __init__(self,systemRole=None, model="glm-4-plus"):
        self.model = model
        self.systemRole = systemRole
        _ = load_dotenv(find_dotenv())
        #self.client = OpenAI()
        self.client = ZhipuAI()

    def execute_task(self, prompt):
        messages = []
        if self.systemRole is not None:
            messages.append({"role":"system", "content":self.systemRole})
        messages.append({"role":"user", "content":prompt}) 

        #print(f"\n\nprompt=\n{prompt}")
        response = self.client.chat.completions.create(
            model=self.model,
            messages=messages
        )
        resultContent = response.choices[0].message.content.strip()
        return resultContent

# 分析缺陷子任务    
class IssueCategory(Task):
    def __init__(self, *args, **kwargs):
        moduleStructure = """
            用户注册与登录
                注册新用户
                用户登录
                密码找回
            商品浏览与搜索
                商品分类浏览
                商品详情查看
                商品搜索
            购物车管理
                添加商品到购物车
                查看购物车内容
                修改购物车商品数量
                移除购物车商品
            订单管理
                创建订单
                查看订单状态
                取消订单
                订单跟踪
            支付与结算
                选择支付方式
                完成支付
                支付状态查询
            用户评价
                查看商品评价
                发表商品评价
            个人中心
                修改个人信息
                查看订单历史
                优惠券管理
                地址管理
        """
        systemRole = f"""你是一名缺陷分析专家，
        你要根据测试人员提交的缺陷说明判断这个缺陷属于系统中的哪一个一级模块。
        当前系统的模块结构如下：
        {moduleStructure}
        你仅需回复一级模块名称即可，
        不要回复推理过程，不要解释，不要其它无关文字。
        """
        super().__init__(systemRole, *args, **kwargs)

    #缺陷模块识别
    def identify(self, issue):
        return super().execute_task(issue)

# 责任人指定子任务
class IssueResponsible(Task):
    def __init__(self,  *args, **kwargs):
        moduleResponsible = """
            用户注册与登录: 张三
            商品浏览与搜索: 李四
            购物车管理: 王五
            订单管理: 赵六
            支付与结算: 陈七
            用户评价: 郑八
            个人中心: 李九
        """
        systemRole = f"""你是一名项目经理，
        你要根据缺陷属于系统中的哪一个模块指派调查修复缺陷的责任人。
        当前系统的模块设计与开发的责任人信息如下：
        {moduleResponsible}
        你仅需回复责任人名字即可，
        不要回复推理过程，不要解释，不要其它无关文字。
        """
        super().__init__(systemRole, *args, **kwargs)

    # 指派责任人
    def assign_responsible(self, category):
        return super().execute_task(category)

def main():

    issue = """
        1.选择一个商品
        2.添加至购物车
        3.继续浏览商品，并选择另一个商品加入购物车
        4.打开购物车，购物车中只有后加入的商品，而没有最初加入的商品
    """

    ic = IssueCategory()
    category = ic.identify(issue)
    print(f"category={category}")

    ir = IssueResponsible()
    responsible = ir.assign_responsible(category)
    print(f"responsible={responsible}")

if __name__ == "__main__":
    main()